from datetime import datetime, date
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
from sklearn.linear_model import LinearRegression
import warnings

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks():
    sql = "SELECT STK_CODE FROM ETF_INDEX_INFO"
    r = exec_sql(sql)
    df1 = pd.DataFrame(r, columns=['代码'])
    df2 = ak.fund_etf_spot_em()
    res = pd.merge(df1, df2, on=['代码'])
    res = res.loc[res['最新价'] < 10, :]
    stocks = res['代码'].to_list()
    # stocks = [tok for tok in stocks if tok != '512870']

    return stocks


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'ETF_DAY_DATA'  # HIS_DAILY_DATA

    res = dict()
    date_data = set_date()
    st = date_data['st']
    et = date_data['et']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
                # sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            else:
                sig_df = df.loc[df['code'] == code]
            sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', drop=False, inplace=True)
            if len(sig_df) <= 220:
                continue
            print(f'code={code}->len={len(sig_df)}')
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            print(f'code={code}->len={len(df)}')
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y-%m-%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


stocks_info = dict()
init_cash = 3e8


def run_backtrade(stocks):
    cerebro = bt.Cerebro()
    cerebro.addstrategy(GridReversionStrategy)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    date_data = set_date()
    st_dt = date_data['st_dt']
    et_dt = date_data['et_dt']

    for i, code in enumerate(data_all):
        if code == index:
            df = data_all[code]
            df = df.loc[:, ['close']]
            df.rename(columns={'close': '399300'}, inplace=True)
            continue
        data = data_all[code]
        print(f"i={i}->len={len(data)}->{code} Done !")
        data_bf = data.loc[data['datetime'] < st_dt]

        if len(data_bf) > 0:
            info_h = data_bf.loc[data_bf['high'].idxmax(), :]
            info_l = data_bf.loc[data_bf['low'].idxmin(), :]
            stocks_info[code] = {'hdt': info_h['datetime'], 'high': info_h['high']}
            stocks_info[code] = {'ldt': info_l['datetime'], 'low': info_l['low']}
        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_dt, todate=et_dt)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    # print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    print(st_dt.date())
    returns = returns[returns.index.date >= st_dt.date()]
    returns.index = returns.index.date
    returns.index = pd.to_datetime(returns.index)

    print(f'returns={returns}')
    html_file = 'test.html'
    qs.reports.html(returns, benchmark=df, title='grid', output=html_file, download_filename=html_file)
    qs.reports.basic(returns)
    f = webbrowser.open(html_file)
    # pf.create_full_tear_sheet(returns)
    # cerebro.plot()


class GridReversionStrategy(bt.Strategy):
    params = (
        ('period', 25),
        ('money', 1e4),
        ('ratio', 0.03),
        ('printlog', False),
        ('maperiod', 26),
        ('times', 1),
        ('span', 220),
        ('st', datetime(2014, 1, 1)),
    )

    def __init__(self):
        self.order = None
        self.daily_pnl = dict()
        self.trade_info = list()
        self.pos_info = dict()
        self.trail_dict = dict()
        self.buy_first = dict()  # 每支ETF是否开始第一天就跌破之前最高价的buy_ratio
        self.cache_log = list()
        self.base_price = dict()  # 每个标的的基准价

        # print(f'stocks_info={stocks_info}')

        stocks = list()
        for data in self.datas:
            code = data._name
            if code != '399300':
                stocks.append(code)
                self.trail_dict[code] = {'hdt': stocks_info[code]['hdt'], 'high': stocks_info[code]['high'],
                                         'ldt': stocks_info[code]['ldt'], 'low': stocks_info[code]['low'],
                                         'cdt': 0, 'cost': 0, 'buy_count': 0, 'pos': dict(), 'grid': dict(), 'counter': 0}

                self.buy_first[code] = True

    def downcast(self, price):
        amount = self.p.money / price // 100 * 100
        return amount

    def set_grid(self, code, price):
        r1, r2 = 1 + self.p.ratio, 1 - self.p.ratio
        self.trail_dict[code]['grid']['lev_0'] = price
        self.trail_dict[code]['grid']['lev_u1'] = price * r1
        self.trail_dict[code]['grid']['lev_u2'] = self.trail_dict[code]['grid']['lev_u1'] * r1
        self.trail_dict[code]['grid']['lev_u3'] = self.trail_dict[code]['grid']['lev_u2'] * r1
        self.trail_dict[code]['grid']['lev_u4'] = self.trail_dict[code]['grid']['lev_u3'] * r1
        self.trail_dict[code]['grid']['lev_u5'] = self.trail_dict[code]['grid']['lev_u4'] * r1
        self.trail_dict[code]['grid']['lev_d1'] = price * r2
        self.trail_dict[code]['grid']['lev_d2'] = self.trail_dict[code]['grid']['lev_d1'] * r2
        self.trail_dict[code]['grid']['lev_d3'] = self.trail_dict[code]['grid']['lev_d2'] * r2
        self.trail_dict[code]['grid']['lev_d4'] = self.trail_dict[code]['grid']['lev_d3'] * r2
        self.trail_dict[code]['grid']['lev_d5'] = self.trail_dict[code]['grid']['lev_d4'] * r2
        self.log(f"lev_0={price}->lev_u1={self.trail_dict[code]['grid']['lev_u1']}->"
                 f"lev_u2={self.trail_dict[code]['grid']['lev_u2']}->lev_u3={self.trail_dict[code]['grid']['lev_u3']}->"
                 f"lev_u4={self.trail_dict[code]['grid']['lev_u4']}->lev_u5={self.trail_dict[code]['grid']['lev_u5']}->"
                 f"lev_d1={self.trail_dict[code]['grid']['lev_d1']}->lev_d2={self.trail_dict[code]['grid']['lev_d2']}->"
                 f"lev_d3={self.trail_dict[code]['grid']['lev_d3']}->lev_d4={self.trail_dict[code]['grid']['lev_d4']}->"
                 f"lev_d5={self.trail_dict[code]['grid']['lev_d5']}")
        self.log(50 * '-')

    def prenext(self):
        self.next()

    def log(self, txt, dt=None, doprint=True):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            content = f'{dt.isoformat()}, {txt}'
            self.cache_log.append(content)
            print(content)

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def next(self):
        if self.order:
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            # dt != cur_dt 防止停牌日触发买卖
            if code == '399300' or dt != cur_dt:
                continue
            elif self.trail_dict[code]['counter'] < self.p.span:
                self.trail_dict[code]['counter'] = self.trail_dict[code]['counter'] + 1
                continue

            if self.buy_first[code]:
                self.set_grid(code, data.lines.close[-1])

            if self.buy_first[code]:
                total_amount = 0
                buy_count = 0

                if self.trail_dict[code]['grid']['lev_d5'] > data.low[0] >= self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.trail_dict[code]['pos']['b1'] = amount
                    self.trail_dict[code]['pos']['b2'] = amount
                    self.trail_dict[code]['pos']['b3'] = amount
                    self.trail_dict[code]['pos']['b4'] = amount
                    self.trail_dict[code]['pos']['b5'] = amount
                    total_amount = amount * 5
                    buy_count = 5
                elif self.trail_dict[code]['grid']['lev_d4'] > data.low[0] >= self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.trail_dict[code]['pos']['b1'] = amount
                    self.trail_dict[code]['pos']['b2'] = amount
                    self.trail_dict[code]['pos']['b3'] = amount
                    self.trail_dict[code]['pos']['b4'] = amount
                    total_amount = amount * 4
                    buy_count = 4
                elif self.trail_dict[code]['grid']['lev_d3'] > data.low[0] >= self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.trail_dict[code]['pos']['b1'] = amount
                    self.trail_dict[code]['pos']['b2'] = amount
                    self.trail_dict[code]['pos']['b3'] = amount
                    total_amount = amount * 3
                    buy_count = 3
                elif self.trail_dict[code]['grid']['lev_d2'] > data.low[0] >= self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.trail_dict[code]['pos']['b1'] = amount
                    self.trail_dict[code]['pos']['b2'] = amount
                    total_amount = amount * 2
                    buy_count = 2
                elif self.trail_dict[code]['grid']['lev_d1'] > data.low[0] >= self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.trail_dict[code]['pos']['b1'] = amount
                    total_amount = amount
                    buy_count = 1
                else:
                    self.buy_first[code] = False

                if total_amount > 0:
                    self.order = self.buy(data=data, size=total_amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]
                    self.trail_dict[code]['buy_count'] = buy_count
            else:
                if self.trail_dict[code]['buy_count'] == 0 and self.trail_dict[code]['grid']['lev_d1'] > data.low[0] >= \
                        self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.log(f'buy lev_d1->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 1 and self.trail_dict[code]['grid']['lev_d2'] > data.low[0] >= \
                        self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.log(f'buy lev_d2->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 2 and self.trail_dict[code]['grid']['lev_d3'] > data.low[0] >= \
                        self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.log(f'buy lev_d3->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 3 and self.trail_dict[code]['grid']['lev_d4'] > data.low[0] >= \
                        self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.log(f'buy lev_d4->xx_a={amount}')
                elif self.trail_dict[code]['buy_count'] == 4 and self.trail_dict[code]['grid']['lev_d5'] > data.low[0] >= \
                        self.trail_dict[code]['low']:
                    amount = self.downcast(data.low[0])
                    self.order = self.buy(data=data, size=amount, name=code)
                    self.log(f'buy lev_d5->xx_a={amount}')

                # 卖出
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 8 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev6']:
                    amount = self.trail_dict[code]['pos'][8]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 7 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev5']:
                    amount = self.trail_dict[code]['pos'][7]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 6 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev4']:
                    amount = self.trail_dict[code]['pos'][6]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 5 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev3']:
                    amount = self.trail_dict[code]['pos'][5]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 4 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev2']:
                    amount = self.trail_dict[code]['pos'][4]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_1'
                    self.log(f'sell lev2->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 3 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev1']:
                    amount = self.trail_dict[code]['pos'][3]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_2'
                    self.log(f'sell lev1->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 2 and \
                        data.high[0] > self.trail_dict[code]['grid']['lev_m']:
                    amount = self.trail_dict[code]['pos'][2]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_3'
                    self.log(f'sell lev_m->in_a={amount}')
                # elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                #         data.high[
                #             0] > self.trail_dict[code]['grid']['lev0'] * 0.85:
                #     amount = self.trail_dict[code]['pos'][1]
                #     self.order = self.sell(data=data, size=amount, name=code)
                #     self.trail_dict[code]['ldt'] = dt
                #     self.trail_dict[code]['low'] = data.high[0]
                #     self.trail_dict[code]['trend'] = 'up1_4'
                #     self.log(f'sell lev0->in_a={amount}')
                elif not self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                        data.close[0] < self.ATR[code] * self.p.times:
                    amount = self.trail_dict[code]['pos'][1]
                    self.order = self.sell(data=data, size=amount, name=code)
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.high[0]
                    self.trail_dict[code]['trend'] = 'up1_4'
                    high_list = data.lines.high.get(ago=-1, size=self.p.span)
                    self.trail_dict[code]['high'] = max(high_list)
                    self.set_grid(code, self.trail_dict[code]['high'])
                    self.log(f'sell lev0->in_a={amount}')
                elif self.trail_dict[code]['trail_high'] and self.trail_dict[code]['buy_count'] == 1 and \
                        data.close[0] < self.trail_dict[code]['high'] * 0.9:
                    self.trail_dict[code]['trail_high'] = False
                    self.order = self.close(data=data, name=code)
                    high_list = data.lines.high.get(ago=-1, size=self.p.span)
                    self.trail_dict[code]['high'] = max(high_list)
                    self.set_grid(code, self.trail_dict[code]['high'])
                    self.log('set_false')
                # elif self.trail_dict[code]['trail_high'] and data.close[0] < self.ATR[code] * self.p.times:
                #     self.trail_dict[code]['trail_high'] = False
                #     self.order = self.close(data=data, name=code)
                #     self.log('set_false')

                if data.low[0] < self.trail_dict[code]['low']:
                    self.trail_dict[code]['ldt'] = dt
                    self.trail_dict[code]['low'] = data.low[0]

                # self.log(f"trail_high={self.trail_dict[code]['trail_high'] }")

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')

                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'BUY', round(order.executed.price, 2), order.executed.size,
                              round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                if self.buy_first[code]:
                    self.buy_first[code] = False
                else:
                    self.trail_dict[code]['buy_count'] += 1
                    self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']] = order.executed.size
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')
                trade_dt = int(bt.num2date(order.created.dt).strftime('%Y%m%d'))
                trade_item = [trade_dt, code, 'SELL', round(order.executed.price, 2),
                              order.executed.size, round(order.executed.comm, 2)]
                self.trade_info.append(trade_item)
                self.trail_dict[code]['cdt'] = int(bt.num2date(order.executed.dt).strftime('%Y%m%d'))
                self.trail_dict[code]['cost'] = round(order.executed.price, 2)
                del self.trail_dict[code]['pos'][self.trail_dict[code]['buy_count']]
                self.trail_dict[code]['buy_count'] -= 1
            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'Order Canceled/Margin/Rejected->code={code}')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def stop(self):
        print('Strategy Finish!!!')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)
        doprint = (False, True)[1]

        if doprint:
            daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
            append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                     columns=['DATE_T', 'PNL'])
            daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
            trade_df = pd.DataFrame(self.trade_info,
                                    columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
            pnl_table = 'TEST_DAILY_PNL'
            trade_table = 'TEST_TRADE_INFO'
            # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
            insert_db(daily_pnl_df, pnl_table)
            insert_db(trade_df, trade_table)
            print(self.trail_dict)
            print('insert table ok')


def set_date():
    st = "20140101"
    et = "20230816"
    st_dt = datetime.strptime(st, '%Y%m%d')
    et_dt = datetime.strptime(et, '%Y%m%d')
    st_d = st_dt.date()
    et_d = et_dt.date()
    res = {'st': st, 'et': et, 'st_d': st_d, 'et_d': et_d, 'st_dt': st_dt, 'et_dt': et_dt}

    return res


def main():
    index = '399300'
    stocks = get_stocks()
    # stocks = ['510880', '159915', '513100', '518880']
    # stocks = ['159901', '510880']
    # stocks = ['159901', '510880', '512510', '512220', '159938', '159909', '513500', '513880', '515770', '515350', '159810', '517350']
    # stocks = stocks[:200]
    # stocks = ['159909', '510170', '159915', '512220', '510880', '510300']
    #
    # stocks = [('159918', '159909', '512220', '159915', '510500', '510510', '159907', '159922', '510290', '159936',
    #            '159930', '159935', '159912', '159903', '159822', '513050', '159902', '510900', '512070', '510410')[19]]
    stocks = ['510900']
    # print(stocks)
    stocks.insert(0, index)
    run_backtrade(stocks)
    # s = get_data(stocks, freq='day', source='not db')
    # print(s)


if __name__ == '__main__':
    main()
